# 8个线性天线3层2次多项式拟合：12天线、sslDb=-30dB、扫描角度0度
import numpy as np
from scipy.signal import windows as SSW # import chebwin
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim.lr_scheduler import ReduceLROnPlateau
from apps.wfs.nlbf_config import NlbfConfig as NG

class ComplexLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super(ComplexLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        
        # 使用两个 nn.Linear 分别处理实部和虚部
        self.linear_real = nn.Linear(in_features, out_features, dtype=torch.float64)
        self.linear_imag = nn.Linear(in_features, out_features, dtype=torch.float64)

    def forward(self, input):
        # 输入的实部和虚部
        input_real = input.real
        input_imag = input.imag
        
        # 分别对实部和虚部进行线性变换
        output_real = self.linear_real(input_real) - self.linear_imag(input_imag)
        output_imag = self.linear_real(input_imag) + self.linear_imag(input_real)
        
        # 合并实部和虚部
        return torch.complex(output_real, output_imag)

class ComplexReLU(torch.nn.Module):
    def __init__(self):
        super(ComplexReLU, self).__init__()
        self.relu = torch.nn.ReLU()

    def forward(self, z):
        # 分别对实部和虚部应用ReLU
        real = self.relu(z.real)
        imag = self.relu(z.imag)
        return torch.complex(real, imag)

class NlbfModel(nn.Module):
    def __init__(self, N_base:int =8, rank:int =3, init_coefs:bool =False, avec_dim:int =1, layers:int = 1):
        super(NlbfModel, self).__init__()
        self.rank = rank
        self.N_base = N_base
        self.layers = layers

        # 波束成形网络
        self.L = [8, 16, 32, 16, 8, 1] # 1361
        # 第1层：从8到16
        self.fc_1 = ComplexLinear(self.L[0], self.L[1])
        self.relu_1 = ComplexReLU()
        #  第2层：从16到32
        self.fc_2 = ComplexLinear(self.L[1], self.L[2])
        self.relu_2 = ComplexReLU()
        # 第3层：从32到16
        self.fc_3 = ComplexLinear(self.L[2], self.L[3])
        self.relu_3 = ComplexReLU()
        # 第4层：从16到8
        self.fc_4 = ComplexLinear(self.L[3], self.L[4])
        self.relu_4 = ComplexReLU()
        # 第5层：从8到1
        self.fc_5 = ComplexLinear(self.L[4], self.L[5])
        self.relu_5 = ComplexReLU()

        # 网络生成网络
        self.ngn_fc1 = ComplexLinear((1000+1)*8, 4096)
        self.ngn_relu1 = ComplexReLU()
        self.ngn_fc2 = ComplexLinear(4096, 2048)
        self.ngn_relu2 = ComplexReLU()
        # self.L = [8, 16, 32, 16, 8, 1] # 1361
        self.ngn_fc3_1_w = ComplexLinear(2048, self.L[0]*self.L[1])
        self.ngn_fc3_1_b = ComplexLinear(2048, self.L[1])
        self.ngn_fc3_2_w = ComplexLinear(2048, self.L[1]*self.L[2])
        self.ngn_fc3_2_b = ComplexLinear(2048, self.L[2])
        self.ngn_fc3_3_w = ComplexLinear(2048, self.L[2]*self.L[3])
        self.ngn_fc3_3_b = ComplexLinear(2048, self.L[3])
        self.ngn_fc3_4_w = ComplexLinear(2048, self.L[3]*self.L[4])
        self.ngn_fc3_4_b = ComplexLinear(2048, self.L[4])
        self.ngn_fc3_5_w = ComplexLinear(2048, self.L[4]*self.L[5])
        self.ngn_fc3_5_b = ComplexLinear(2048, self.L[5])

        
            
    
    def forward(self, X:torch.Tensor, scan_v:torch.Tensor, A:torch.Tensor) -> torch.Tensor:
        '''
        仅讨论8天线情形
        输入参数：
            X: ((1000,8), complex128) 目标信号+干扰信号+噪声
            scan_v: ((8,), complex128) 扫描方向
            A: ((1800,8), complex128) 导向向量
        输出：((1800,1), float64)
        '''
        a1 = self.relu_1(self.fc_1(A))
        a2 = self.relu_2(self.fc_2(a1))
        a3 = self.relu_3(self.fc_3(a2))
        a4 = self.relu_4(self.fc_4(a3))
        a5 = self.fc_5(a4)
        outputs = torch.abs(a5) / torch.max(torch.abs(a5))
        return outputs

    def save(self, pt_fn='./work/wfs/nlbf.pt') -> None:
        torch.save(self.state_dict(), pt_fn)

    def load(self, pt_fn='./work/wfs/nlbf.pt') -> None:
        self.load_state_dict(torch.load(pt_fn))

    def get_weights(self) -> torch.tensor:
        return self.coefficients[0]     